import math
import matplotlib.pyplot as plt

# ===================== 距离函数 =====================
def euclidean_distance(x1, x2):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(x1, x2)))

# ===================== KNN 预测函数 =====================
def knn_predict(train_X, train_y, x, k=3):
    distances = []
    for xi, yi in zip(train_X, train_y):
        d = euclidean_distance(x, xi)
        distances.append((d, yi, xi))
    
    # 按距离排序
    distances.sort(key=lambda t: t[0])
    k_nearest = distances[:k]

    # 多数投票
    labels = [label for _, label, _ in k_nearest]
    pred = max(set(labels), key=labels.count)
    return pred, k_nearest

# ===================== 示例数据 =====================
train_X = [
    [1.0, 2.0],
    [1.5, 1.8],
    [5.0, 8.0],
    [6.0, 9.0]
]
train_y = [0, 0, 1, 1]

test_point = [2.0, 2.0]
k = 3

# 预测
pred_label, nearest = knn_predict(train_X, train_y, test_point, k)
print("预测类别为：", pred_label)

# ===================== 可视化 =====================
plt.figure(figsize=(7, 7))

# 颜色选择
colors = {0: "blue", 1: "red"}

# 画训练点
for (x, y), label in zip(train_X, train_y):
    plt.scatter(x, y, color=colors[label], s=80)
    plt.text(x+0.05, y+0.05, f"{label}", fontsize=12)

# 画测试点
plt.scatter(test_point[0], test_point[1], color="green", marker="^", s=200)
plt.text(test_point[0]+0.05, test_point[1]+0.05, "test", fontsize=12)

# 画 k 个最近邻连线
for dist, label, point in nearest:
    plt.plot([test_point[0], point[0]], [test_point[1], point[1]], linestyle="--")

plt.title(f"KNN (k={k}) 预测结果: {pred_label}")
plt.xlabel("X1")
plt.ylabel("X2")
plt.grid(True)
plt.show()
